#!/usr/bin/env python
"""
Simple command line tool to extract sequences of interest from a sequence file.
e.g.
python seqDB.py -db infile.sto -format stockholm -accs "acc1" "acc2" -desc "Methanobacter maripuladis" "Methanococcus"
"""
# Copyright (c) Alex Leach. 2011
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>
#
# Description:
# Program to reformat or filter sequences in a variety of formats.
# See http://www.biopython.org/wiki/SeqIO for a list of formats, 
# although there are currently more supported by biopython.
# e.g. sff format (Roche 454's binary sequence format) is supported
# by biopython, but not mentioned on the webpage.

from Bio import SeqIO
import sys,re


class Options():
    options = {      '-db' : None ,
                     '-accs': [] ,
                     '-desc' : [] ,
                     '-format' : 'fasta',
                     '-out' : sys.stdout,
                     '-outformat' : 'fasta',
                     '--reverse' : False
            }

    helpTxt = {     '-db' : 'File which contains sequences of interest',
                    '-accs' : 'List of accessions to look for within <-db>. This is relaxed, so returns all matches which contain a given accession, even if it isnt the complete accession. Separate with spaces.',
                    '-desc' : 'Descriptions to look for within db. Again, is relaxed, so returns all matches. Envelope each description with " and separate with spaces.',
                    '-format' : 'The input sequence format. Default: fasta',
                    '-out' : 'Output file. Default: standard out. N.B. Appends to the file.',
                    '-outformat' : 'Output file format. Default: fasta.',
                    '--reverse' : 'Reverse transcribe the sequences'
              }
    def __getitem__( self, key ):
        return self.options[key]
    def __setitem__( self, key, value ):
        self.options.update( { key : value } )
        return 
    def printHelp(self):
            print 'Valid options:-'
            for key,val in self.helpTxt.items():
                    print key,' : ',val

    def parseArgs(self, args):
            if len(args) == 0:
                    self.printHelp()
                    exit()
            prevargDash = False
            for i in xrange(len(args)):
                    if args[i].startswith('--'):
                        other = self.options[args[i]]
                        self.options[args[i]] = not other
                    elif args[i].startswith('-'):
                            if args[i] in self.options:
                                    prevargDash = True
                                    command = args[i]
                            else:
                                    print 'Invalid option: {0}'.format( args[i] )
                                    self.printHelp()
                                    exit()
                    elif prevargDash:
                            if command in ['-db','-outformat','-out']:
                                    self.options.update( { command :  args[i] } )
                                    prevargDash = False
                            elif command == '-accs':
                                    self.options[command].append(args[i])
                            elif command == '-desc':
                                    self.options[command].append( '('+args[i] + ')')
                            elif command == '-format':
                                    self.options.update( {command : args[i] } )
                                    prevargDash = False
                    else:
                            print "Don't know what to do with {0}".format( args[i] )
                            exit()
            return self.options

class Writer():
    @staticmethod
    def reverser(seq,form):
        seq.seq = seq.seq.reverse_complement()
        return seq.format(form)
    @staticmethod
    def writer(seq,form):
        return seq.format(form)

def main(options):
    if len(options['-desc']) > 0:
            desc_reg = re.compile( '|'.join( options['-desc'] ), re.I )
    if len(options['-accs']) > 0:
            acc_reg = re.compile( '|'.join( options['-accs'] ), re.I )
    if options['--reverse']:
        writer = Writer().reverser
    else:
        writer = Writer().writer
    Ndesc = len(options['-desc'])
    Naccs = len(options['-accs'])
    printed = False
    if options['-out'] != sys.stdout:
        out = file( options['-out'] , 'w' )
    else:
        out = sys.stdout
    with file(options['-db'],'r') as infile:
        count = 0
        for seq in SeqIO.parse( infile,options['-format']  ):
            count += 1
            if Naccs > 0:
                accFind = acc_reg.search( seq.id )
                if accFind:
                    out.write( writer(seq,options['-outformat']) )
                    #out.write( seq.seq.tostring().upper())
                    printed = True
            elif Ndesc > 0:
                desc_find = desc_reg.search( seq.description )
                if desc_find:
                    printed = True
                    out.write( writer(seq, options['-outformat'] ) )
            else:
                printed = True
                out.write( writer(seq,options['-outformat'] ) )
    if not printed and count > 0:
        sys.stderr.write( "No matching sequences found. Here's the last sequence header as an example...\n" )
        sys.stderr.write( seq.format( options['-outformat']).splitlines()[0]  )
        sys.stderr.write('\n')
    else:
        out.write( '\n')
    if options['-out'] != sys.stdin:
        options['-out'].close()

if __name__ == '__main__':
        args = sys.argv[1:]
        options = Options()
        options.parseArgs(args)
        main(options)

